package com.example;

import java.net.InetAddress;
import java.net.NetworkInterface;
import java.net.SocketException;
import java.net.UnknownHostException;
import java.nio.ByteBuffer;
import java.util.Arrays;

/**
 * 对一个以太帧的简单包装，提供了一些方便的方法来读取包的内容。
 */
public final class Packet {
    private final byte[] packet;

    /**
     * Packet类的构造函数。
     * @param packet 表示数据包的字节数组。
     */
    Packet(byte[] packet) {
        this.packet = packet;
    }

    /**
     * 检查TCP头部的SYN标志是否被设置。
     * @return 如果SYN标志被设置，返回true，否则返回false。
     */
    public boolean isSYN() {
        return (packet[47] & 0x02) != 0;
    }

    /**
     * 检查TCP头部的ACK标志是否被设置。
     * @return 如果ACK标志被设置，返回true，否则返回false。
     */
    public boolean isACK() {
        return (packet[47] & 0x10) != 0;
    }

    /**
     * 检查TCP头部的FIN标志是否被设置。
     * @return 如果FIN标志被设置，返回true，否则返回false。
     */
    public boolean isFIN() {
        return (packet[47] & 0x01) != 0;
    }

    /**
     * 从TCP头部获取窗口大小。
     * @return 作为整数的窗口大小。
     */
    public int windowSize() {
        int tcpHeaderIndex = 14 + 20; // 以太网帧头部 + IP包头
        return (packet[tcpHeaderIndex + 14] & 0xFF) << 8 | (packet[tcpHeaderIndex + 15] & 0xFF);
    }

    /**
     * 从TCP头部获取序列号。
     * @return 作为整数的序列号。
     */
    public int seqNumber() {
        int tcpHeaderIndex = 14 + 20; // 以太网帧头部 + IP包头
        return (packet[tcpHeaderIndex + 4] & 0xFF) << 24 | (packet[tcpHeaderIndex + 5] & 0xFF) << 16 | (packet[tcpHeaderIndex + 6] & 0xFF) << 8 | (packet[tcpHeaderIndex + 7] & 0xFF);
    }

    /**
     * 从TCP头部获取确认号。
     * @return 作为整数的确认号。
     */
    public int ackNumber() {
        int tcpHeaderIndex = 14 + 20; // 以太网帧头部 + IP包头
        return (packet[tcpHeaderIndex + 8] & 0xFF) << 24 | (packet[tcpHeaderIndex + 9] & 0xFF) << 16 | (packet[tcpHeaderIndex + 10] & 0xFF) << 8 | (packet[tcpHeaderIndex + 11] & 0xFF);
    }

    /**
     * 提取并返回数据包中的数据，不包括头部。
     * @return 作为字节数组的数据。
     */
    public byte[] data() {
        int tcpHeaderIndex = 14 + 20; // 以太网帧头部 + IP包头
        int headerSize = ((packet[tcpHeaderIndex + 12] >> 4) & 0x0F) * 4; // TCP头部大小
//        System.out.println("packet.length: " + packet.length);
//        System.out.println("tcpHeaderSize: " + headerSize);
        byte[] data = new byte[packet.length - tcpHeaderIndex - headerSize]; // 以太网帧头部 + IP包头 + TCP头部
//        System.out.println("data.length: " + data.length);
        System.arraycopy(packet, tcpHeaderIndex + headerSize, data, 0, data.length);
        return data;
    }

    /**
     * 返回表示数据包的字节数组。
     * @return 作为字节数组的数据包。
     */
    public byte[] bytes() {
        return packet;
    }

    /**
     * 创建一个新的PacketBuilder。
     * @return 一个新的PacketBuilder。
     */
    public static PacketBuilder builder() {
        return new PacketBuilder();
    }

    /**
     * 使用指定的数据包类型创建一个新的PacketBuilder。
     * @param type 数据包类型。
     * @return 一个新的PacketBuilder。
     */
    public static PacketBuilder builder(PktType type) {
        return new PacketBuilder(type);
    }

    /**
     * 使用指定的数据、目标IP、源端口、目标端口、序列号和确认号创建一个新的PacketBuilder。
     * @param data 数据。
     * @param desIP 目标IP。
     * @param srcPort 源端口。
     * @param desPort 目标端口。
     * @param seq 序列号。
     * @param ack 确认号。
     * @return 一个新的PacketBuilder。
     */
    public static PacketBuilder builder(byte[] data, byte[] desIP, short srcPort, short desPort, int seq, int ack) {
        return new PacketBuilder(data, desIP, srcPort, desPort, seq, ack);
    }

    /**
     * 使用指定的数据包类型、目标IP、源端口、目标端口、序列号和确认号创建一个新的PacketBuilder。
     * @param type 数据包类型。
     * @param desIP 目标IP。
     * @param srcPort 源端口。
     * @param desPort 目标端口。
     * @param seq 序列号。
     * @param ack 确认号。
     * @return 一个新的PacketBuilder。
     */
    public static PacketBuilder builder(PktType type, byte[] desIP, short srcPort, short desPort, int seq, int ack) {
        return new PacketBuilder(type, desIP, srcPort, desPort, seq, ack);
    }

    /**
     * 从字节数组构建一个Packet。
     * @param packet 字节数组。
     * @return 一个新的Packet。
     */
    public static Packet buildFromBytes(byte[] packet) {
        // 从ip头中获得ip数据包总长度
        int ethernetHeaderLength = 14; // 以太网帧头部
        int totalLength = (packet[ethernetHeaderLength + 2] & 0xFF) << 8 | (packet[ethernetHeaderLength + 3] & 0xFF);

        // 去除padding
        if (packet.length > ethernetHeaderLength + totalLength) {
            return new Packet(Arrays.copyOf(packet, ethernetHeaderLength + totalLength));
        } else {
            return new Packet(packet);
        }

    }

    /**
     * 表示可能用到的标志位的枚举。包括表示发送数据，不设置标志位的NONE。
     */
    public enum PktType {
        SYN(0x02),
        ACK(0x10),
        FIN(0x01),
        SYN_ACK(0x12),
        FIN_ACK(0x11),
        NONE(0x00);

        private final byte flags;

        /**
         * PktType枚举的构造函数。
         * @param flags 标志位。
         */
        PktType(int flags) {
            this.flags = (byte) flags;
        }

        /**
         * 返回标志位。
         * @return 作为字节的标志位。
         */
        public byte getFlags() {
            return flags;
        }
    }

    /**
     * 计算校验和。
     * @param bytes 字节数组。
     * @return 作为短整数的校验和。
     */
    private static short checksum(byte[] bytes) {
        // 确认输入的字节数组不是null且长度不为0
        if (null == bytes || bytes.length == 0) {
            throw new IllegalArgumentException("bytes to checksum is null or empty");
        }

        int sum = 0;

        // 添加字节
        byte[] paddedBytes = new byte[(bytes.length % 2 == 0) ? bytes.length : bytes.length + 1];
        System.arraycopy(bytes, 0, paddedBytes, 0, bytes.length);
        for (int i = 0; i < paddedBytes.length; i += 2) {
            int word = ((paddedBytes[i] & 0xFF) << 8) | (paddedBytes[i + 1] & 0xFF);
            sum += word;
        }

        // 添加溢出
        while ((sum >> 16) > 0) {
            sum = (sum & 0xFFFF) + (sum >> 16);
        }

        // 返回补码
        return (short) ~sum;
    }


    /**
     * 用于构造数据包的类。
     * <p>
     * 使用方法类似如下：
     * <pre>
     * {@code
     * Packet packet = Packet.builder()
     *     .setSYN()
     *     .setACK()
     *     .setWindowSize(100)
     *     .setData(new byte[] {1, 2, 3})
     *     .build();
     * </pre>
     */
    public static class PacketBuilder {

        private ByteBuffer ethernetFrame;

        private ByteBuffer ipPacket;

        private ByteBuffer tcpPacket;

        private byte[] data = new byte[0];

        PacketBuilder() {
            this(PktType.NONE, new byte[4], (short) 0, (short) 0, 0, 0);
        }

        PacketBuilder(PktType type) {
            this(type, new byte[4], (short) 0, (short) 0, 0, 0);
        }

        // 构造包含数据的TCP包。
        PacketBuilder(byte[] data, byte[] desIP, short srcPort, short desPort, int seq, int ack) {
            buildHeader(PktType.NONE, desIP, srcPort, desPort, seq, ack);

            // 设置数据
            this.data = data;
        }

        // 构造不包含数据的TCP包。
        PacketBuilder(PktType type, byte[] desIP, short srcPort, short desPort, int seq, int ack) {
            buildHeader(type, desIP, srcPort, desPort, seq, ack);
        }

        private void buildHeader(PktType type, byte[] desIp, short srcPort, short desPort, int seq, int ack) {
            // 获取网关MAC地址、本机MAC地址、本机IP地址
            byte[] gatewayMac;
            byte[] localMac;
            byte[] localIp;
            try {
                gatewayMac = Utils.getGatewayMac();
                var localHost = InetAddress.getLocalHost();
                localIp = localHost.getAddress();
                localMac = NetworkInterface.getByInetAddress(localHost).getHardwareAddress();

            } catch (UnknownHostException e) {
                throw new RuntimeException("Failed to get local host address", e);
            } catch (SocketException e) {
                throw new RuntimeException("Failed to get local network interface", e);
            } catch (Throwable e) {
                throw new RuntimeException("Failed to get gateway MAC address", e);
            }

            ByteBuffer ethernetFrame = ByteBuffer.allocate(14); // 以太网帧头
            ethernetFrame.put(gatewayMac); // 目的MAC地址
            ethernetFrame.put(localMac); // 源MAC地址
            ethernetFrame.put((byte) 0x08); // 类型，0x08表示IP
            ethernetFrame.put((byte) 0x00);
            this.ethernetFrame = ethernetFrame;

            ByteBuffer ipPacket = ByteBuffer.allocate(20); // IP包头
            ipPacket.put((byte) 0x45); // 版本和头部长度
            ipPacket.put((byte) 0x00); // 服务类型
            ipPacket.putShort((short) 0); // 总长度，稍后填充
            ipPacket.putShort((short) 0); // 重组标识。
            ipPacket.putShort((short) 0); // 标志和片偏移。
            ipPacket.put((byte) 0x40); // TTL，0x40表示64
            ipPacket.put((byte) 0x06); // 协议，0x06表示TCP
            ipPacket.putShort((short) 0); // 头部校验和，稍后计算
            ipPacket.put(localIp); // 源IP地址
            ipPacket.put(desIp); // 目标IP地址
            this.ipPacket = ipPacket;

            ByteBuffer tcpPacket = ByteBuffer.allocate(20); // TCP包头
            tcpPacket.putShort(srcPort); // 源端口
            tcpPacket.putShort(desPort); // 目标端口
            tcpPacket.putInt(seq); // 序列号
            tcpPacket.putInt(ack); // 确认号
            tcpPacket.put((byte) 0x50); // 数据偏移，0x50表示头部长度为5（20字节）
            tcpPacket.put(type.getFlags()); // 标志位
            tcpPacket.putShort((short) 65535); // 窗口大小，默认为65535
            tcpPacket.putShort((short) 0); // 校验和，稍后计算
            tcpPacket.putShort((short) 0); // 紧急指针，不使用
            this.tcpPacket = tcpPacket;

        }


        public PacketBuilder setSYN() {
            byte flag = tcpPacket.get(13);
            tcpPacket.put(13, (byte) (flag | 0x02));
            return this;
        }

        public PacketBuilder setACK() {
            byte flag = tcpPacket.get(13);
            tcpPacket.put(13, (byte) (flag | 0x10));
            return this;
        }

        public PacketBuilder setFIN() {
            byte flag = tcpPacket.get(13);
            tcpPacket.put(13, (byte) (flag | 0x01));
            return this;
        }

        public PacketBuilder setAckNumber(int ack) {
            tcpPacket.putInt(8, ack);
            return this;
        }

        public PacketBuilder setSeqNumber(int seq) {
            tcpPacket.putInt(4, seq);
            return this;
        }

        public PacketBuilder setWindowSize(int size) {
            tcpPacket.putShort(14, (short) size);
            return this;
        }

        public PacketBuilder setData(byte[] data) {
            this.data = data;
            return this;
        }

        public PacketBuilder setSourcePort(short port) {
            tcpPacket.putShort(0, port);
            return this;
        }

        public PacketBuilder setDestinationPort(short port) {
            tcpPacket.putShort(2, port);
            return this;
        }

        public PacketBuilder setDestinationIp(byte[] ip) {
            ipPacket.position(16);
            ipPacket.put(ip);
            return this;
        }

        public Packet build() {
            // 填充IP包头的总长度
            ipPacket.putShort(2, (short) (ipPacket.capacity() + tcpPacket.capacity() + data.length));

            // 填充IP包头的校验和
            ipPacket.putShort(10, checksum(ipPacket.array()));
            byte[] ipHeader = ipPacket.array();

            // 构造TCP伪首部
            byte[] pseudoHeader = new byte[12];
            System.arraycopy(ipHeader, 12, pseudoHeader, 0, 4);
            System.arraycopy(ipHeader, 16, pseudoHeader, 4, 4);
            // 第八字节为0
            pseudoHeader[9] = (byte) 0x06; // 协议号，0x06代表TCP
            pseudoHeader[10] = (byte) ((tcpPacket.capacity() + data.length >> 8) & 0xFF);
            pseudoHeader[11] = (byte) (tcpPacket.capacity() + data.length & 0xFF);

            // 计算TCP校验和
            byte[] sumArray = new byte[pseudoHeader.length + tcpPacket.capacity() + data.length];
            System.arraycopy(pseudoHeader, 0, sumArray, 0, pseudoHeader.length);
            System.arraycopy(tcpPacket.array(), 0, sumArray, pseudoHeader.length, tcpPacket.capacity());
            System.arraycopy(data, 0, sumArray, pseudoHeader.length + tcpPacket.capacity(), data.length);
            tcpPacket.putShort(16, checksum(sumArray));

            // 检查数据包长度，如果不足60字节，添加padding
            int totalLength = ethernetFrame.capacity() + ipPacket.capacity() + tcpPacket.capacity() + data.length;
            byte[] padding = new byte[0];
            if (totalLength < 60) {
                padding = new byte[60 - totalLength];
                Arrays.fill(padding, (byte) 0);
            }

            // 构造完整数据包
            ByteBuffer packet = ByteBuffer.allocate(ethernetFrame.capacity() + ipPacket.capacity() + tcpPacket.capacity() + data.length + padding.length);
            packet.put(ethernetFrame.array());
            packet.put(ipHeader);
            packet.put(tcpPacket.array());
            packet.put(data);
            packet.put(padding);
            return new Packet(packet.array());
        }
    }
}

